package com.xzz.multipleMerge;

import java.util.*;

/**
 * @author: hhz
 * @create: 2022-01-14 10:22
 * 373. 查找和最小的K对数字
 * 给定两个以升序排列的整数数组 nums1 和 nums2,以及一个整数 k。
 * <p>
 * 定义一对值(u,v)，其中第一个元素来自nums1，第二个元素来自 nums2。
 * <p>
 * 请找到和最小的 k个数对(u1,v1), (u2,v2) ... (uk,vk)。
 **/
public class KSmallestPairs {
    public static void main(String[] args) {
        System.out.println(kSmallestPairs(new int[]{1, 7, 11}, new int[]{2, 4, 6}, 3));
        System.out.println(kSmallestPairs(new int[]{1, 1, 2}, new int[]{1, 2, 3}, 2));
        System.out.println(kSmallestPairs(new int[]{1, 2}, new int[]{3}, 3));
    }

    public static List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new ArrayList<>();
        //小顶堆---这里学到了多路归并的思路
        PriorityQueue<List<Integer>> queue = new PriorityQueue<>(
                (o1, o2) -> nums1[o1.get(0)] + nums2[o1.get(1)] - nums1[o2.get(0)] - nums2[o2.get(1)]
        );
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            queue.offer(Arrays.asList(i, 0));
        }
        while (res.size() < k && !queue.isEmpty()) {
            List<Integer> poll = queue.poll();
            int i = poll.get(0);
            int j = poll.get(1);
            res.add(Arrays.asList(nums1[i], nums2[j]));
            if (j + 1 < nums2.length) {
                queue.offer(Arrays.asList(i, j + 1));
            }
        }
        return res;
    }
}
